import os.path as osp
import os
import sys
import itertools

# sys.path.append(osp.join(osp.dirname(__file__), "..", ".."))
import cv2
import numpy as np

from src.dust3r.datasets.base.base_multiview_dataset import BaseMultiViewDataset, cropping
from src.dust3r.utils.image import imread_cv2

# Depth images represent Euclidean distances in meters from the camera center.
# Euclidean distances (in meters) to the optical center of the camera
class HyperSim_Multi(BaseMultiViewDataset):
    def __init__(self, *args, ROOT, samples_per_scene=10, max_interval=16, min_interval=1, **kwargs):
        self.video = True
        self.is_metric = True
        self.max_interval = max_interval
        self.min_interval = min_interval
        self.samples_per_scene = samples_per_scene
        # self.split = split
        super().__init__(*args, **kwargs)

        # self.ROOT = osp.join(ROOT, self.split)
        self.ROOT = ROOT
        self.loaded_data = self._load_data()

    def _load_data(self):
        self.all_scenes = sorted(
            [f for f in os.listdir(self.ROOT) if os.path.isdir(osp.join(self.ROOT, f))]
        )
        subscenes = []
        for scene in self.all_scenes:
            # not empty
            subscenes.extend(
                [
                    osp.join(scene, f)
                    for f in os.listdir(osp.join(self.ROOT, scene))
                    if os.path.isdir(osp.join(self.ROOT, scene, f))
                    and len(os.listdir(osp.join(self.ROOT, scene, f))) > 0
                ]
            )

        offset = 0
        scenes = []
        images = []
        scene_img_list = []
        for scene_idx, scene in enumerate(subscenes):
            scene_dir = osp.join(self.ROOT, scene)
            rgb_paths = sorted([f for f in os.listdir(scene_dir) if f.endswith(".png")])
            assert len(rgb_paths) > 0, f"{scene_dir} is empty."
            num_imgs = len(rgb_paths)
            cut_off = (
                self.num_views*self.min_interval if not self.allow_repeat else max(self.num_views*self.min_interval // 3, 3)
            )
            if num_imgs < cut_off:
                print(f"Skipping {scene}")
                continue
            img_ids = list(np.arange(num_imgs) + offset)

            scenes.append(scene)
            scene_img_list.append(img_ids)
            images.extend(rgb_paths)
            offset += num_imgs

        self.scenes = scenes
        self.images = images
        self.scene_img_list = scene_img_list

    def __len__(self):
        return len(self.scenes) * self.samples_per_scene

    def get_image_num(self):
        return len(self.images)

    def _get_views(self, scene_id, resolution, rng, num_views):
        scene_id = scene_id // self.samples_per_scene
        all_image_ids = self.scene_img_list[scene_id]

        try:
            pos, ordered_video = self.efficient_random_intervals_revised(
                0,
                len(all_image_ids),
                num_views,
                rng,
                min_interval=self.min_interval,
                max_interval=self.max_interval,
            )
        except ValueError as e:
            print(f"Error in _get_views of {scene_id}, {num_views} of {len(all_image_ids)}: {e}")
            return None

        image_idxs = np.array(all_image_ids)[pos]
        views = []
        for v, view_idx in enumerate(image_idxs):
            scene_dir = osp.join(self.ROOT, self.scenes[scene_id])

            rgb_path = self.images[view_idx]
            depth_path = rgb_path.replace("rgb.png", "depth.npy")
            cam_path = rgb_path.replace("rgb.png", "cam.npz")

            rgb_image = imread_cv2(osp.join(scene_dir, rgb_path), cv2.IMREAD_COLOR)
            depthmap = np.load(osp.join(scene_dir, depth_path)).astype(np.float32)
            depthmap[~np.isfinite(depthmap)] = 0  # invalid
            cam_file = np.load(osp.join(scene_dir, cam_path))
            intrinsics = cam_file["intrinsics"].astype(np.float32)
            camera_pose = cam_file["pose"].astype(np.float32)

            rgb_image, depthmap, intrinsics = self._crop_resize_if_necessary(
                rgb_image, depthmap, intrinsics, resolution, rng=rng, info=view_idx
            )
            intrinsics, camera_pose = cropping.get_center_camera(intrinsics, camera_pose, depthmap=depthmap)
            # generate img mask and raymap mask
            img_mask, ray_mask = self.get_img_and_ray_masks(
                self.is_metric, v, rng, p=[0.75, 0.2, 0.05]
            )

            views.append(
                dict(
                    img=rgb_image,
                    depthmap=depthmap.astype(np.float32),
                    camera_pose=camera_pose.astype(np.float32),
                    camera_intrinsics=intrinsics.astype(np.float32),
                    dataset="hypersim",
                    label=self.scenes[scene_id] + "-" + rgb_path,
                    instance=f"{str(scene_id)}_{str(view_idx)}",
                    is_metric=self.is_metric,
                    is_video=ordered_video,
                    quantile=np.array(1.0, dtype=np.float32),
                    img_mask=img_mask,
                    ray_mask=ray_mask,
                    camera_only=False,
                    depth_only=False,
                    single_view=False,
                    reset=False,
                )
            )
        assert len(views) == num_views
        return views

# if __name__ == "__main__":
#     from dust3r.datasets.utils.transforms import ImgNorm, SeqColorJitter
#     dataset = HyperSim_Multi(allow_repeat=False, split='train', ROOT="/lc/data/3D/processed_hypersim", aug_crop=1, resolution=224, transform=SeqColorJitter, num_views=4, n_corres=0)
#     print(type(dataset), len(dataset))
#     print(dataset.get_image_num())

#     rng = np.random.default_rng()
#     views = dataset._get_views(2, (320, 228), rng, 1)
#     print(len(views), type(views[0]))
#     view = views[0]
#     print(view.keys())
#     img, depth,pose = view['img'], view['depthmap'], view['camera_pose']
#     print(depth.shape, pose.shape)
#     print(type(img))